clear;

% This program solves the stochastic Ramsey growth model using the
% ET-algorithm by Den Haan, Kobielarz, and Rendahl (2015).

% Declare parameters for the model

alpha = 1/3;        % Capital share of income: 1/3.
beta = 1.04^(-1/4); % Discount factor: 4 percent annual real interest rate.
delta = 0.025;      % Depreciation rate.
gamma = 10;         % Coefficient of relative risk aversion.
sigma = 0.03;       % Standard deviation of productivity shock.
rho = 0.95;         % Persistence of technology shock.

% Implied steady state
kss = ((1/beta+delta-1)/alpha)^(1/(alpha-1));
css = kss^(alpha)-delta*kss;

% Information regarding the shock
N_quad = 5;                 % Number of nodes for the quadrature (5 goes a long way).
[Z W] = hernodes(N_quad);   % Generate nodes and weights.
W = pi^(-1/2)*W;            % Normalize weights.
Z = Z*sqrt(2)*sigma;        % Normalize nodes.

% Declar symbolic variables

syms k c z phi_k phi_z kt ct zt

kp = exp(z)*k^(alpha)+(1-delta)*k-c;
zp = rho*z+Z;

% Define the Euler Equation

EE = c-(W'*( beta*(1+exp(zp)*alpha*kp^(alpha-1)-delta).*(ct+phi_k*(kp-kt)+phi_z*(zp-zt)).^(-gamma) ))^(-1/gamma);

% Take derivatives

dk = diff(EE,k);
dz = diff(EE,z);
dc = diff(EE,c);

% The implicit function theorem reveals that the following equations should
% hold

Ek = phi_k+dk./dc;
Ez = phi_z+dz./dc;

% The collection of equations to be solved is therefore

E = [EE;Ek;Ez];
E = subs(E,[kt zt ct],[k z c]);

% Convert into a matlab function

E = matlabFunction(E,'vars',{[c;phi_k;phi_z],k,z});

% Solve the problem along a stochastic simulation.

T = 500;

k0 = zeros(1,T+1);
z0 = zeros(1,T+1);

k0(1) = kss;
z0(1) = 0;

% Initial guess
X = [css;0.9;0.9];

options = optimset('Display','off','TolFun',1e-14,'TolX',1e-14);
rng(20150215,'twister');

e = sigma*randn(T,1);

tic
for t = 1:T
    disp(t)
    E1 = @(x) E(x,k0(t),z0(t));
    X(:,t+1) = fsolve(E1,X(:,t),options);
    k0(t+1) = exp(z0(t))*k0(t)^(alpha)+(1-delta)*k0(t)-X(1,t+1);
    z0(t+1) = rho*z0(t)+e(t);
    
end
toc

c0 = X(1,:);

% Done.

